import collections
import math
import os
import shutil
import pandas as pd
import torch
import torchvision
from torch import nn
from d2l import torch as d2l
from torch.utils.data import Dataset,DataLoader
from PIL import Image

data_dir = ('/root/autodl-tmp/leave')

def read_csv_labels(fname):
    with open(fname,'r') as f:
        lines = f.readlines()[1:]
        tokens = [l.rstrip().split(',') for l in lines]
        return dict(((name,label) for name,label in tokens))

labels = read_csv_labels(os.path.join(data_dir,'train.csv'))
num_classes = len(collections.Counter(labels.values()))


train_ends = len(labels)-1
test_begin = len(labels)


def copyfile(filename,target_dir):
    os.makedirs(target_dir,exist_ok=True)
    shutil.copy(filename,target_dir)

def create_dir(target_dir,file_dir,begin,end):
    for i in range(begin,end+1):
        copyfile(os.path.join(file_dir,f'{i}.jpg'),target_dir)


create_dir(os.path.join(data_dir,'train'),os.path.join(data_dir,'images'),0,train_ends)
create_dir(os.path.join(data_dir,'test'),os.path.join(data_dir,'images'),test_begin,27152)

def reorg_train(data_dir,labels):
    for train_file in os.listdir(os.path.join(data_dir,'train')):
        label = labels['images/'+train_file]
        fname = os.path.join(data_dir,'train',train_file)
        copyfile(fname,os.path.join(data_dir,'train_common','train',label))

#@save
def reorg_test(data_dir):
    """在预测期间整理测试集，以方便读取"""
    for test_file in os.listdir(os.path.join(data_dir, 'test')):
        copyfile(os.path.join(data_dir, 'test', test_file),
                 os.path.join(data_dir, 'train_common', 'test',
                              'unknown'))

reorg_train(data_dir,labels)
reorg_test(data_dir)


transform_train = torchvision.transforms.Compose([
    # 在高度和宽度上将图像放大到40像素的正方形
    torchvision.transforms.Resize(300),
    # 随机裁剪出一个高度和宽度均为40像素的正方形图像，
    # 生成一个面积为原始图像面积0.64～1倍的小正方形，
    # 然后将其缩放为高度和宽度均为32像素的正方形
    torchvision.transforms.RandomResizedCrop(224, scale=(0.64, 1.0),
                                                   ratio=(1.0, 1.0)),
    torchvision.transforms.RandomHorizontalFlip(),
    torchvision.transforms.ToTensor(),
    # 标准化图像的每个通道
    torchvision.transforms.Normalize([0.4914, 0.4822, 0.4465],
                                     [0.2023, 0.1994, 0.2010])])

transform_test = torchvision.transforms.Compose([
    torchvision.transforms.ToTensor(),
    torchvision.transforms.Normalize([0.4914, 0.4822, 0.4465],
                                     [0.2023, 0.1994, 0.2010])])


train_ds = torchvision.datasets.ImageFolder(
    root= '/root/autodl-tmp/leave/train_common/train',
    transform=transform_train)

test_ds = torchvision.datasets.ImageFolder(
   root= '/root/autodl-tmp/leave/train_common/test',
    transform=transform_test )

batch_size = 128
train_iter = torch.utils.data.DataLoader(
    dataset=train_ds, batch_size=batch_size, shuffle=True, drop_last=True)



test_iter = torch.utils.data.DataLoader(dataset=test_ds, batch_size=batch_size, shuffle=False,
                                        drop_last=False)

def get_net():
    net = torchvision.models.resnet18()
    num_in = net.fc.in_features
    net.fc = nn.Linear(num_in, num_classes)
    return net

loss = nn.CrossEntropyLoss(reduction="none")
net = get_net()


lr = 0.01
trainer = torch.optim.Adam(lr=lr,params=net.parameters())

d2l.train_ch13(net,train_iter,test_iter,loss,trainer,10)
d2l.plt.show()